package com.baidu.mpks.statistic.elasticSearch;

import com.baidu.mpks.statistic.meta.Action;
import lombok.extern.slf4j.Slf4j;
import org.elasticsearch.action.get.GetRequest;
import org.elasticsearch.action.search.MultiSearchRequest;
import org.elasticsearch.action.search.MultiSearchResponse;
import org.elasticsearch.action.search.SearchRequest;
import org.elasticsearch.action.search.SearchResponse;
import org.elasticsearch.action.update.UpdateRequest;
import org.elasticsearch.client.RequestOptions;
import org.elasticsearch.client.RestHighLevelClient;
import org.elasticsearch.index.query.*;
import org.elasticsearch.script.Script;
import org.elasticsearch.script.ScriptType;
import org.elasticsearch.search.aggregations.AggregationBuilders;
import org.elasticsearch.search.aggregations.PipelineAggregationBuilder;
import org.elasticsearch.search.aggregations.PipelineAggregatorBuilders;
import org.elasticsearch.search.aggregations.bucket.histogram.DateHistogramAggregationBuilder;
import org.elasticsearch.search.aggregations.bucket.histogram.DateHistogramInterval;
import org.elasticsearch.search.aggregations.bucket.range.RangeAggregationBuilder;
import org.elasticsearch.search.aggregations.bucket.terms.TermsAggregationBuilder;
import org.elasticsearch.search.aggregations.metrics.CardinalityAggregationBuilder;
import org.elasticsearch.search.builder.SearchSourceBuilder;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.stereotype.Service;
import org.springframework.util.StringUtils;

import javax.annotation.Resource;
import java.io.IOException;
import java.time.ZoneId;
import java.util.Arrays;
import java.util.HashMap;
import java.util.Objects;
import java.util.stream.Collectors;

@Slf4j
@Service
public class EsBaseQuery {
    @Resource
    private RestHighLevelClient esClient;

    @Value("${es.index}")
    private String statisticIndex;

    public static final long SEARCH_TIMEOUT = 60;
    public static final ZoneId zoneId = ZoneId.of("Asia/Shanghai");
    public static final String TS_FIELD = "tstamp";

    /**
     * 返回指定日期区间内的，指定字段值的总的文档数
     *
     * @param filed
     * @param value
     * @param from
     * @param to
     * @return
     * @throws IOException
     */
    public SearchResponse countRange(String filed, String value, double from, double to) throws IOException {
        SearchRequest searchRequest = initSearchRequest();
        SearchSourceBuilder searchSourceBuilder = new SearchSourceBuilder();
        searchSourceBuilder.size(0);
        RangeAggregationBuilder rangeAgg = AggregationBuilders.range("range")
            .field(TS_FIELD).addRange("date_range", from, to);
        CardinalityAggregationBuilder cardinalityAgg = AggregationBuilders.cardinality("cardinality")
                .field(EsService.FIELD_LOGID);
        rangeAgg.subAggregation(cardinalityAgg);
        searchSourceBuilder.aggregation(rangeAgg).query(QueryBuilders.termQuery(filed, value));
        searchRequest.source(searchSourceBuilder);

        log.info("[countRange] search request {}", searchRequest.toString());

        return esClient.search(searchRequest, RequestOptions.DEFAULT);
    }

    /**
     * 返回指定日期区间内的，限定条件下的指定字段值的总的文档数
     *
     * @param filed
     * @param value
     * @param from
     * @param to
     * @return
     * @throws IOException
     */
    public SearchResponse countRange(String filed, String value, double from, double to,
            BoolQueryBuilder boolQueryBuilder) throws IOException {
        SearchRequest searchRequest = initSearchRequest();
        SearchSourceBuilder searchSourceBuilder = new SearchSourceBuilder();
        searchSourceBuilder.size(0);
        RangeAggregationBuilder rangeAgg = AggregationBuilders.range("range")
            .field(TS_FIELD).addRange("date_range", from, to);
        if (Objects.equals(value, Action.LOGIN.toString())) {
            rangeAgg.subAggregation(
                    AggregationBuilders.cardinality("cardinality")
                            .field(EsService.FIELD_USERID));
        }

        boolQueryBuilder.must(QueryBuilders.termQuery(filed, value));
        searchSourceBuilder.aggregation(rangeAgg).query(boolQueryBuilder);
        searchRequest.source(searchSourceBuilder);

        log.info("[countRange] search request {}", searchRequest.toString());

        return esClient.search(searchRequest, RequestOptions.DEFAULT);
    }

    /**
     * 根据特定条件返回指定日期区间的文档数量，以interval为维度，可去重
     *
     * @param filed
     * @param value
     * @param cardinalities
     *            去重字段
     * @param from
     * @param to
     * @param interval
     *            时间间隔，支持hour、day、week和month
     * @param boolQueryBuilder
     * @return
     * @throws IOException
     */
    public SearchResponse countHistogram(String filed, String value, String[] cardinalities, double from, double to,
            String interval, BoolQueryBuilder boolQueryBuilder) throws IOException {
        SearchRequest searchRequest = initSearchRequest();
        SearchSourceBuilder searchSourceBuilder = new SearchSourceBuilder();
        searchSourceBuilder.size(0);
        RangeAggregationBuilder rangeAgg = AggregationBuilders.range("range")
            .field(TS_FIELD).addRange("date_range", from, to);
        String actualInterval;
        switch (interval) {
            case "hour" :
                actualInterval = DateHistogramInterval.HOUR.toString();
                break;
            case "day" :
                actualInterval = DateHistogramInterval.DAY.toString();
                break;
            case "week" :
                actualInterval = DateHistogramInterval.WEEK.toString();
                break;
            case "month" :
                actualInterval = DateHistogramInterval.MONTH.toString();
                break;
            default :
                actualInterval = DateHistogramInterval.DAY.toString();
        }
        DateHistogramInterval dateHistogramInterval = new DateHistogramInterval(actualInterval);

        DateHistogramAggregationBuilder histoAgg = AggregationBuilders.dateHistogram("histogram")
            .field(TS_FIELD).calendarInterval(dateHistogramInterval).timeZone(zoneId).minDocCount(1);
        if (cardinalities != null && cardinalities.length > 0) {
            CardinalityAggregationBuilder cardinalityAgg = cardinalities.length > 1
                    ? getCardinalityAgg(cardinalities)
                    : AggregationBuilders.cardinality("cardinality").field(cardinalities[0]);

            histoAgg.subAggregation(cardinalityAgg);
        }
        rangeAgg.subAggregation(histoAgg);
        TermQueryBuilder termQueryBuilder = QueryBuilders.termQuery(filed, value);
        if (boolQueryBuilder != null) {
            boolQueryBuilder.must(termQueryBuilder);
            searchSourceBuilder.query(boolQueryBuilder);
        } else {
            searchSourceBuilder.query(termQueryBuilder);
        }
        searchSourceBuilder.aggregation(rangeAgg);
        searchRequest.source(searchSourceBuilder);

        log.info("[countHistogram] search request {}", searchRequest.toString());

        return esClient.search(searchRequest, RequestOptions.DEFAULT);
    }

    public MultiSearchResponse multiCountRange(String filed, String value, double from, double to,
            QueryBuilder... queryBuilders) throws IOException {
        MultiSearchRequest multiSearchRequest = new MultiSearchRequest();
        RangeAggregationBuilder rangeAgg = AggregationBuilders.range("range")
            .field(TS_FIELD).addRange("date_range", from, to);

        for (QueryBuilder queryBuilder : queryBuilders) {
            SearchRequest searchRequest = initSearchRequest();
            SearchSourceBuilder searchSourceBuilder = new SearchSourceBuilder();
            searchSourceBuilder.aggregation(rangeAgg).query(QueryBuilders.termQuery(filed, value)).query(queryBuilder);
            searchRequest.source(searchSourceBuilder);
            multiSearchRequest.add(searchRequest);
        }

        return esClient.msearch(multiSearchRequest, RequestOptions.DEFAULT);
    }

    /**
     * 返回指定日期区间限定条件下的某个值出现次数
     *
     * @param service
     * @param actionField
     * @param aggField
     * @param from
     * @param to
     * @param size
     * @param queryBuilder
     * @return
     * @throws IOException
     */
    public SearchResponse aggregationTerm(String service, String actionField, String aggField, String cardinality,
            double from, double to, int size, QueryBuilder queryBuilder) throws IOException {
        SearchRequest searchRequest = initSearchRequest();
        SearchSourceBuilder searchSourceBuilder = new SearchSourceBuilder();
        searchSourceBuilder.size(0);

        TermsAggregationBuilder termAgg = AggregationBuilders.terms("count").field(aggField).size(size);
        CardinalityAggregationBuilder cardinalityAgg = AggregationBuilders.cardinality("cardinality")
                .field(cardinality);
        termAgg.subAggregation(cardinalityAgg);

        RangeAggregationBuilder rangeAgg = AggregationBuilders.range("range")
            .field(TS_FIELD).addRange("date_range", from, to);
        rangeAgg.subAggregation(termAgg);
        searchSourceBuilder.aggregation(rangeAgg);
        BoolQueryBuilder termQueryBuilder = QueryBuilders.boolQuery()
                .must(QueryBuilders.termQuery(service, actionField));
        if (queryBuilder != null) {
            termQueryBuilder.must(queryBuilder);
        }
        searchSourceBuilder.query(termQueryBuilder);
        searchRequest.source(searchSourceBuilder);

        log.info("[aggregationTerm] search request {}", searchRequest.toString());

        return esClient.search(searchRequest, RequestOptions.DEFAULT);
    }

    /**
     * 返回指定日期区间限定条件下的某个值出现次数，可去重
     *
     * @param service
     * @param actionFields
     * @param scriptCode
     * @param from
     * @param to
     * @param cardinalityField
     * @param cardCondition
     *            可进一步对去重的结果进行过滤
     * @return
     * @throws Exception
     */
    public SearchResponse aggregationTerms(String service, String[] actionFields, String scriptCode,
            String cardinalityField, String cardCondition, double from, double to, int size) throws Exception {
        SearchRequest searchRequest = initSearchRequest();
        SearchSourceBuilder searchSourceBuilder = new SearchSourceBuilder();
        searchSourceBuilder.size(0);

        TermsQueryBuilder termQueryBuilder = QueryBuilders.termsQuery(service, actionFields);
        // 没有lodId的不进行无点击query统计
        ExistsQueryBuilder existsQueryBuilder = new ExistsQueryBuilder("logId.keyword");
        BoolQueryBuilder boolQueryBuilder = new BoolQueryBuilder();
        boolQueryBuilder.filter(termQueryBuilder);
        boolQueryBuilder.must(existsQueryBuilder);
        searchSourceBuilder.query(boolQueryBuilder);

        CardinalityAggregationBuilder cardinalityAgg = AggregationBuilders.cardinality("cardinality")
                .field(cardinalityField);
        PipelineAggregationBuilder filterAgg = null;
        if (!StringUtils.isEmpty(cardCondition)) {
            Script filtersScript = new Script(ScriptType.INLINE, "painless", cardCondition, new HashMap<>());
            filterAgg = PipelineAggregatorBuilders
                .bucketSelector("select", filtersScript, "cardinality");

        }

        TermsAggregationBuilder termsAgg = AggregationBuilders.terms("terms");
        termsAgg.size(size);

        Script termScript = new Script(ScriptType.INLINE, "painless", scriptCode, new HashMap<>());
        termsAgg.script(termScript);
        termsAgg.subAggregation(cardinalityAgg);
        if (filterAgg != null) {
            termsAgg.subAggregation(filterAgg);
        }

        RangeAggregationBuilder rangeAgg = AggregationBuilders.range("range")
            .field(TS_FIELD).addRange("date_range", from, to);
        rangeAgg.subAggregation(termsAgg);

        searchSourceBuilder.aggregation(rangeAgg);
        searchRequest.source(searchSourceBuilder);

        log.info("[aggregationTerms] search request {}", searchRequest.toString());

        return esClient.search(searchRequest, RequestOptions.DEFAULT);
    }

    /**
     * 根据特定条件返回指定日期区间的文档数量，可去重
     *
     * @param service
     * @param actionField
     * @param cardinalities
     * @param from
     * @param to
     * @param queryBuilder
     * @return
     * @throws IOException
     */
    public SearchResponse aggregationCardinality(String service, String actionField, String[] cardinalities,
            double from, double to, QueryBuilder queryBuilder) throws IOException {
        SearchRequest searchRequest = initSearchRequest();
        SearchSourceBuilder searchSourceBuilder = new SearchSourceBuilder();
        searchSourceBuilder.size(0);

        CardinalityAggregationBuilder cardinalityAgg = cardinalities.length > 1
                ? getCardinalityAgg(cardinalities)
                : AggregationBuilders.cardinality("cardinality").field(cardinalities[0]);

        RangeAggregationBuilder rangeAgg = AggregationBuilders.range("range")
            .field(TS_FIELD).addRange("date_range", from, to);
        rangeAgg.subAggregation(cardinalityAgg);

        BoolQueryBuilder termQueryBuilder = QueryBuilders.boolQuery()
                .must(QueryBuilders.termQuery(service, actionField));
        if (null != queryBuilder) {
            termQueryBuilder.must(queryBuilder);
        }
        searchSourceBuilder.aggregation(rangeAgg).query(termQueryBuilder);
        searchRequest.source(searchSourceBuilder);

        log.info("[aggregationCardinality] search request {}", searchRequest.toString());

        return esClient.search(searchRequest, RequestOptions.DEFAULT);
    }

    public TermsAggregationBuilder getTermAggRecursive(String[] terms, int index) throws Exception {
        TermsAggregationBuilder termAgg = null;
        if (terms.length > index) {
            String term = terms[index];
            termAgg = AggregationBuilders.terms(terms[index]).field(term).size(Integer.MAX_VALUE - 1);
        }
        if (terms.length - 1 == index) {
            return termAgg;
        }
        TermsAggregationBuilder subAgg = getTermAggRecursive(terms, index + 1);
        if (subAgg != null) {
            Objects.requireNonNull(termAgg).subAggregation(subAgg);
        }
        return termAgg;
    }

    public CardinalityAggregationBuilder getCardinalityAgg(String[] cardinalities) {
        CardinalityAggregationBuilder cardinalityAgg = AggregationBuilders.cardinality("cardinality");
        String scriptCode = Arrays.stream(cardinalities).map(e -> "doc['" + e + "'].value")
                .collect(Collectors.joining("+'_'+"));
        Script script = new Script(ScriptType.INLINE, "painless", scriptCode, new HashMap<>());
        cardinalityAgg.script(script);
        return cardinalityAgg;
    }

    public SearchResponse searchByRequest(SearchRequest searchRequest) throws IOException {
        searchRequest.indices(statisticIndex);

        log.info("[searchByRequest] search request {}", searchRequest.toString());

        return esClient.search(searchRequest, RequestOptions.DEFAULT);
    }

    public GetRequest initGetRequest(String id) {
        return new GetRequest(statisticIndex, id);
    }

    public SearchRequest initSearchRequest() {
        return new SearchRequest(statisticIndex);
    }

    public UpdateRequest initUpdateRequest(String id) {
        return new UpdateRequest(statisticIndex, id);
    }

    public void close() throws IOException {
        if (esClient != null) {
            esClient.close();
        }

    }

}
